#!/bin/bash

# AI Booster 远程部署脚本
# 用于将安装脚本部署到远程Linux服务器

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m'

print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示帮助信息
show_help() {
    echo "AI Booster 远程部署脚本"
    echo ""
    echo "使用方法:"
    echo "  $0 <server-ip> [username] [port]"
    echo ""
    echo "参数说明:"
    echo "  server-ip  - 目标服务器IP地址"
    echo "  username   - SSH用户名 (默认: root)"
    echo "  port       - SSH端口 (默认: 22)"
    echo ""
    echo "示例:"
    echo "  $0 192.168.1.100"
    echo "  $0 192.168.1.100 root 22"
    echo "  $0 192.168.1.100 ubuntu 2222"
    echo ""
    echo "注意:"
    echo "  - 需要配置SSH密钥认证或准备输入密码"
    echo "  - 确保目标服务器安装了bash"
    echo "  - 建议使用root用户或具有sudo权限的用户"
}

# 检查必需文件
check_files() {
    local missing_files=()
    
    if [ ! -f "install_ai_booster.sh" ]; then
        missing_files+=("install_ai_booster.sh")
    fi
    
    if [ ! -f "manage_ai_booster.sh" ]; then
        missing_files+=("manage_ai_booster.sh")
    fi
    
    if [ ${#missing_files[@]} -gt 0 ]; then
        print_error "缺少必需文件:"
        for file in "${missing_files[@]}"; do
            echo "  - $file"
        done
        echo ""
        print_info "请确保在包含安装脚本的目录中运行此脚本"
        exit 1
    fi
    
    print_success "所有必需文件已找到"
}

# 部署到远程服务器
deploy_to_server() {
    local server_ip="$1"
    local username="${2:-root}"
    local port="${3:-22}"
    
    print_info "开始部署到服务器: $username@$server_ip:$port"
    
    # 检查SSH连接
    if ! ssh -p "$port" -o ConnectTimeout=10 -o BatchMode=yes "$username@$server_ip" "echo 'SSH连接成功'" 2>/dev/null; then
        print_warning "SSH密钥认证失败，将使用密码认证"
        print_info "请准备输入SSH密码"
    fi
    
    # 上传文件
    print_info "正在上传安装脚本..."
    if scp -P "$port" install_ai_booster.sh "$username@$server_ip:/tmp/"; then
        print_success "install_ai_booster.sh 上传成功"
    else
        print_error "install_ai_booster.sh 上传失败"
        exit 1
    fi
    
    if scp -P "$port" manage_ai_booster.sh "$username@$server_ip:/tmp/"; then
        print_success "manage_ai_booster.sh 上传成功"
    else
        print_error "manage_ai_booster.sh 上传失败"
        exit 1
    fi
    
    # 上传README（可选）
    if [ -f "README.md" ]; then
        scp -P "$port" README.md "$username@$server_ip:/tmp/" 2>/dev/null && \
        print_info "README.md 上传成功"
    fi
    
    # 在远程服务器上设置权限并提供安装指令
    print_info "正在远程服务器上设置权限..."
    
    ssh -p "$port" "$username@$server_ip" << 'EOF'
cd /tmp
chmod +x install_ai_booster.sh
chmod +x manage_ai_booster.sh
echo ""
echo "========================================="
echo "     文件上传完成！"
echo "========================================="
echo ""
echo "文件位置: /tmp/"
echo "  - install_ai_booster.sh  (安装脚本)"
echo "  - manage_ai_booster.sh   (管理脚本)"
echo "  - README.md             (说明文档)"
echo ""
echo "安装命令:"
echo "  sudo bash /tmp/install_ai_booster.sh"
echo ""
echo "管理命令:"
echo "  sudo bash /tmp/manage_ai_booster.sh config"
echo ""
echo "========================================="
EOF
    
    if [ $? -eq 0 ]; then
        print_success "远程部署完成！"
        echo ""
        print_info "下一步操作:"
        echo "  1. SSH登录到服务器: ssh -p $port $username@$server_ip"
        echo "  2. 运行安装命令: sudo bash /tmp/install_ai_booster.sh"
        echo "  3. 开始使用: source /etc/ai_booster enable"
    else
        print_error "远程命令执行失败"
        exit 1
    fi
}

# 批量部署
batch_deploy() {
    local servers_file="$1"
    
    if [ ! -f "$servers_file" ]; then
        print_error "服务器列表文件不存在: $servers_file"
        exit 1
    fi
    
    print_info "开始批量部署，读取服务器列表: $servers_file"
    print_warning "文件格式: ip:username:port (每行一个服务器)"
    
    local line_num=0
    local success_count=0
    local fail_count=0
    
    while IFS=':' read -r ip username port; do
        line_num=$((line_num + 1))
        
        # 跳过空行和注释
        if [[ -z "$ip" || "$ip" =~ ^[[:space:]]*# ]]; then
            continue
        fi
        
        # 设置默认值
        username="${username:-root}"
        port="${port:-22}"
        
        echo ""
        echo "========================================="
        echo "部署到服务器 $line_num: $ip"
        echo "========================================="
        
        if deploy_to_server "$ip" "$username" "$port"; then
            success_count=$((success_count + 1))
            print_success "服务器 $ip 部署成功"
        else
            fail_count=$((fail_count + 1))
            print_error "服务器 $ip 部署失败"
        fi
        
        sleep 2  # 避免连接过快
    done < "$servers_file"
    
    echo ""
    echo "========================================="
    echo "批量部署完成！"
    echo "========================================="
    echo "成功: $success_count 台服务器"
    echo "失败: $fail_count 台服务器"
    echo "总计: $((success_count + fail_count)) 台服务器"
}

# 创建服务器列表模板
create_server_list_template() {
    cat > servers.txt << 'EOF'
# AI Booster 服务器列表
# 格式: IP:用户名:端口
# 示例:
# 192.168.1.100:root:22
# 192.168.1.101:ubuntu:2222
# 192.168.1.102::   (使用默认用户名root和端口22)

EOF
    
    print_success "服务器列表模板已创建: servers.txt"
    print_info "请编辑 servers.txt 文件，添加您的服务器信息"
}

# 主函数
main() {
    case "$1" in
        "help"|"-h"|"--help"|"")
            show_help
            ;;
        "batch")
            check_files
            batch_deploy "${2:-servers.txt}"
            ;;
        "create-template")
            create_server_list_template
            ;;
        *)
            if [[ "$1" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
                check_files
                deploy_to_server "$1" "$2" "$3"
            else
                print_error "无效的IP地址: $1"
                echo ""
                show_help
                exit 1
            fi
            ;;
    esac
}

# 运行主函数
main "$@"
